{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a LangGraph Agent\n",
    "\n",
    "The following example demonstrates building an agent using LangGraph. Zep is used to personalize agent responses based on information learned from prior conversations. \n",
    "\n",
    "The agent implements:\n",
    "- persistance of new chat turns to Zep and recall of relevant Facts using the most recent messages.\n",
    "- an in-memory MemorySaver to maintain agent state. We use this to add recent chat history to the agent prompt. As an alternative, you could use Zep for this. \n",
    "\n",
    "**IMPRTANT**: You should consider truncating MemorySaver's chat history as by default LangGraph state grows unbounded. We've included this in our example below. See the LangGraph documentation for insight.\n",
    "\n",
    "## Install dependencies\n",
    "```shell\n",
    "pip install zep-cloud langchain-openai langgraph ipywidgets\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import uuid\n",
    "from contextlib import suppress\n",
    "from typing import Annotated\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import Image, display\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_logging():\n",
    "    logger = logging.getLogger()\n",
    "    logger.setLevel(logging.ERROR)\n",
    "    console_handler = logging.StreamHandler(sys.stdout)\n",
    "    console_handler.setLevel(logging.ERROR)\n",
    "    formatter = logging.Formatter(\"%(name)s - %(levelname)s - %(message)s\")\n",
    "    console_handler.setFormatter(formatter)\n",
    "    logger.addHandler(console_handler)\n",
    "    return logger\n",
    "\n",
    "\n",
    "logger = setup_logging()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangSmith integration (Optional)\n",
    "\n",
    "If you'd like to trace your agent using LangSmith, ensure that you have a `LANGSMITH_API_KEY` set in your environment.\n",
    "\n",
    "Then set `os.environ['LANGCHAIN_TRACING_V2'] = 'false'` to `true`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"false\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Zep LangGraph Tutorial\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Zep\n",
    "\n",
    "Ensure that you've configured the following API key in your environment. We're using Zep's Async client here, but we could also use the non-async equivalent.\n",
    "\n",
    "```bash\n",
    "ZEP_API_KEY=\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zep_cloud import Message\n",
    "from zep_cloud.client import AsyncZep\n",
    "\n",
    "zep = AsyncZep(api_key=os.environ.get(\"ZEP_API_KEY\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage, SystemMessage, trim_messages\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import END, START, StateGraph, add_messages\n",
    "from langgraph.prebuilt import ToolNode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Zep's Search as a Tool\n",
    "These are examples of simple Tools that searches Zep for facts (from edges) or nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "    first_name: str\n",
    "    last_name: str\n",
    "    session_id: str\n",
    "    user_name: str\n",
    "\n",
    "\n",
    "@tool\n",
    "async def search_facts(state: State, query: str, limit: int = 5) -> list[str]:\n",
    "    \"\"\"Search for facts in all conversations had with a user.\n",
    "\n",
    "    Args:\n",
    "        state (State): The Agent's state.\n",
    "        query (str): The search query.\n",
    "        limit (int): The number of results to return. Defaults to 5.\n",
    "\n",
    "    Returns:\n",
    "        list: A list of facts that match the search query.\n",
    "    \"\"\"\n",
    "    edges = await zep.graph.search(\n",
    "        user_id=state[\"user_name\"], text=query, limit=limit, search_scope=\"edges\"\n",
    "    )\n",
    "    return [edge.fact for edge in edges]\n",
    "\n",
    "\n",
    "@tool\n",
    "async def search_nodes(state: State, query: str, limit: int = 5) -> list[str]:\n",
    "    \"\"\"Search for nodes in all conversations had with a user.\n",
    "\n",
    "    Args:\n",
    "        state (State): The Agent's state.\n",
    "        query (str): The search query.\n",
    "        limit (int): The number of results to return. Defaults to 5.\n",
    "\n",
    "    Returns:\n",
    "        list: A list of node summaries for nodes that match the search query.\n",
    "    \"\"\"\n",
    "    nodes = await zep.graph.search(\n",
    "        user_id=state[\"user_name\"], text=query, limit=limit, search_scope=\"nodes\"\n",
    "    )\n",
    "    return [node.summary for node in nodes]\n",
    "\n",
    "\n",
    "tools = [search_facts, search_nodes]\n",
    "\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0).bind_tools(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chatbot Function Explanation\n",
    "\n",
    "The chatbot uses Zep to provide context-aware responses. Here's how it works:\n",
    "\n",
    "1. **Context Retrieval**: It retrieves relevant facts for the user's current conversation (session). Zep uses the most recent messages to determine what facts to retrieve.\n",
    "\n",
    "2. **System Message**: It constructs a system message incorporating the facts retrieved in 1., setting the context for the AI's response.\n",
    "\n",
    "3. **Message Persistence**: After generating a response, it asynchronously adds the user and assistant messages to Zep. New Facts are created and existing Facts updated using this new information.\n",
    "\n",
    "4. **Messages in State**: We use LangGraph state to store the most recent messages and add these to the Agent prompt. We limit the message list to the most recent 3 messages for demonstration purposes. You may also use Zep's chat history for this purpose. \n",
    "\n",
    "This approach enables the chatbot to maintain context across interactions and provide personalized responses based on the user's history and preferences stored in Zep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def chatbot(state: State):\n",
    "    memory = await zep.memory.get(state[\"session_id\"])\n",
    "\n",
    "    system_message = SystemMessage(\n",
    "        content=f\"\"\"You are a compassionate mental health bot and caregiver. Review information about the user and their prior conversation below and respond accordingly.\n",
    "        Keep responses empathetic and supportive. And remember, always prioritize the user's well-being and mental health.\n",
    "\n",
    "        {memory.context}\"\"\"\n",
    "    )\n",
    "\n",
    "    messages = [system_message] + state[\"messages\"]\n",
    "\n",
    "    response = await llm.ainvoke(messages)\n",
    "\n",
    "    # Add the new chat turn to the Zep graph\n",
    "    messages_to_save = [\n",
    "        Message(\n",
    "            role_type=\"user\",\n",
    "            role=state[\"first_name\"] + \" \" + state[\"last_name\"],\n",
    "            content=state[\"messages\"][-1].content,\n",
    "        ),\n",
    "        Message(role_type=\"assistant\", content=response.content),\n",
    "    ]\n",
    "\n",
    "    await zep.memory.add(\n",
    "        session_id=state[\"session_id\"],\n",
    "        messages=messages_to_save,\n",
    "    )\n",
    "\n",
    "    # Truncate the chat history to keep the state from growing unbounded\n",
    "    # In this example, we going to keep the state small for demonstration purposes\n",
    "    # We'll use Zep's Facts to maintain conversation context\n",
    "    state[\"messages\"] = trim_messages(\n",
    "        state[\"messages\"],\n",
    "        strategy=\"last\",\n",
    "        token_counter=len,\n",
    "        max_tokens=3,\n",
    "        start_on=\"human\",\n",
    "        end_on=(\"human\", \"tool\"),\n",
    "        include_system=True,\n",
    "    )\n",
    "\n",
    "    logger.info(f\"Messages in state: {state['messages']}\")\n",
    "\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Agent\n",
    "\n",
    "This section sets up the Agent's LangGraph graph:\n",
    "\n",
    "1. **Graph Structure**: It defines a graph with nodes for the agent (chatbot) and tools, connected in a loop.\n",
    "\n",
    "2. **Conditional Logic**: The `should_continue` function determines whether to end the graph execution or continue to the tools node based on the presence of tool calls.\n",
    "\n",
    "3. **Memory Management**: It uses a MemorySaver to maintain conversation state across turns. This is in addition to using Zep for facts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_builder = StateGraph(State)\n",
    "\n",
    "memory = MemorySaver()\n",
    "\n",
    "\n",
    "# Define the function that determines whether to continue or not\n",
    "async def should_continue(state, config):\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    # If there is no function call, then we finish\n",
    "    if not last_message.tool_calls:\n",
    "        return \"end\"\n",
    "    # Otherwise if there is, we continue\n",
    "    else:\n",
    "        return \"continue\"\n",
    "\n",
    "\n",
    "graph_builder.add_node(\"agent\", chatbot)\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "graph_builder.add_edge(START, \"agent\")\n",
    "\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"agent\", should_continue, {\"continue\": \"tools\", \"end\": END}\n",
    ")\n",
    "graph_builder.add_edge(\"tools\", \"agent\")\n",
    "\n",
    "\n",
    "graph = graph_builder.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our LangGraph agent graph is illustrated below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPYAAAERCAIAAADHRs0RAAAAAXNSR0IArs4c6QAAIABJREFUeJztnXdcU9f//08G2QsIQaYMFZFNlTpw4KobRKutC1erftRq6x7tF62r7roqatXWtuJWHHVrFReK4BZEEWQGQkhCErJ/f1x/lCIjam7uzc15Pvwjud6c87rJi3PPPee834dkMpkABEJcyFgLgEDQBVocQnCgxSEEB1ocQnCgxSEEB1ocQnCoWAvAjLI3GqVcr5QbdBqjRm3EWo5Z0JlkqgOJxaNyBFSRFx1rObYByd7GxXMfK189UuY+rvIOZGvVBjaf6iii6bQ2YnEGpaJUq5TrKVRS3lOlbzDHL5TTIoyNtS5cY0cWf/mg6sZJibs/w8OP6RvMYbBtu5Om05hyH1flP1fnZyk7DhC2juJirQin2IXFtWrj+T9KKFRSx4FCvtABazkWRiU33DxVLhXrPhvtynMm2tV9PMS3eNFL9aldxfHTPIUeNKy1oIisXHdie2F0rItfCOy3/AeCW7yiRHv1kDh+uifWQqzE6V+LI2Mc3fwYWAvBEUS2+KtHysyrlfHTPbAWYlVO7Sz2DWYHdeBhLQQv2PYjVyPIJbrUE2X25m8AwICv3J7ekZe8rsZaCF4grMUvHxCPmt8caxXY8PlMzztnK7TVtjESijbEtPjtMxKPFkwylYS1EMxoEca5caIcaxW4gIAW12lMD/6pbNfbCWshWBLUgffmhUou0WEtBHsIaPGMK9Kun4usU5fBYMjMzPzgj1dVVT1//tyiiv6lc5zo4XUZSoXbEAS0+JNbMq+WLOvU9eOPP65YseKDP/7FF1+cOHHCoor+pXkg6+H1SpQKtyGIZnHxGw2LR2XzKdapTqPRfNgHkbFarVZraUX/QqYArwBW3jMVelXYBEQbF79/SUqmkMK7CSxecmpq6ubNmwsKCtzd3YcOHTp8+PDExMRTp07VnJCSkuLu7p6SknLw4MGcnBwWi9WhQ4fZs2c7OjoCAC5evDh//vy1a9fu27fvyZMnCQkJp0+fLikpQT7brFmz2kVZiqx7CkmxtuNAZ4uXbEMQbTFtWZHGN8jyM9gqlWrevHl+fn6LFy/OyckpKysDAIwfP760tLSwsHDp0qUAAKFQCAB49OiRj49Pv379KioqkpOTlUrlxo0ba8r56aefpk6dOmXKFG9v765du06bNu2TTz4ZOXIkjYbK4gI2n/r0jhyNkm0IollcKTOweZa/qIqKCo1G07179759+9Yc9Pb2FggEEokkPDy85uDChQtJpLeDlVQqdffu3RqNhk5/u7Z7+PDhAwYMQF6LRCIqlSoUCmt/3LKw+VSlTI9S4bYC0SyukuvZfMtflIeHR2ho6K+//spkMuPj4xtpdHU6XXJy8pkzZ0pKShgMhtFolEqlzZo1Q/43KirK4toagc2jKOUGa9aIQ4j2uEmlkSlky8/4kEikTZs2DRgwYOPGjfHx8ffv36/3NJPJNHPmzN27dw8aNGjLli39+vUDABiN/84yslhWGupBIFNINAbRfuL3hWjX70AjVclRuTVzOJz58+cfOXKEw+F89913KtXbkYraz+v3799PS0ubP3/+iBEjgoODW7Ro0WSxqD7uK2V6ih1P8SIQzeIsHlWJjsWR8UEPD48vvviiqqqqqKgIAMBkMiUSSU07XVlZCQBo3bp17be1W/E6MJnM8nIUp9lVcgObZ6XxU9xCtL64iwddi0KssU6nGzJkSK9evfz9/Q8dOsThcDw9PQEAkZGRKSkpK1asCA8P5/F4ISEhNBpty5YtgwcPfvHixZ49ewAAOTk5yMnvEhERcfbs2b179/J4vNDQUHNa/fdCrTQ0a27va8cpiYmJWGuwJCYTeHhNFhhl4dXSSqUyPz//ypUrly9fdnFxSUxMRFzbokULmUx29uzZ+/fvCwSCmJgYPz+/kydPnjx5Uq/XL1u2TCwWZ2ZmDhgw4NWrVxcvXhw2bJhA8O+YfWhoaFZW1pkzZ54/fx4UFOTr62tZ2WlnKzxbsZzdiBzu1CREm/oBACTNezl+qZ8D3d77oACAX+a8/HqFH8XBrr8KonVUAABBHfhvslWNhDBevHhx2bJl7x6n0+kNTcjv2bPH4k1sHaqqqmqGzOvg6OgolUrfPb5x48ZGxtSLXlYHtOXaub+J2YrLynUpSUWjFzUYD6FWq+t1jFarbWjAG5mmsajMuhiNxpr5/DrodDoHh3pC652dnWsmld7l8M8F0bHCZj723hcnYCvOFzp4tGA+uS0Pal9/j5zJZDKZTKvragIymezu7m6p0nKfKBlsCvQ3AQcNEToOEL56UIW1CizJvlfVcYAQaxW4gJgWZ7DJYd0EJ7YXYS0EGy78Wdo8iOXUDKYNAoS1OADAO4DlHcC6tF+MtRBrk3qinM2jtm4L87+9hYCPm7V5+VCZ91TZ/Qsrxblhzo0UiUDoENQRJlH5F8K24gj+oWwXT/qRTQUGPZH/khFO7yqmM8nQ33UgeCuOUJxbffVQmX8oO6oPMcPyM65UZlyRxgwT+QbDhIZ1sQuLIxP7d89VpF+StvvMybsVS+RNhPzz5UXavKfKjCvSwE95HfoLyfa+4Kp+7MXiCAad6cE1Wc4DRVWlPvBTvsloYvEoPCcHo9E2vgQKhSSv0KkUBpMRvMhQ0Bhk/zBOaLTA1nOlo4p9WbwGldxQmKOWS3UquQEAU5XMwqExRUVFOp2ueXMLZ5zjCKgmo4nNo3IEVHc/JteJgDN3FsdOvyMWj9IykoNe+fv2XVBIJH3HfopeFRAzgTc4CMGBFocQHGhxVGAymVwunF/EBdDiqKBWqxUKBdYqIABaHC2oVCra68shZgItjgp6vV6vt/c0VDgBWhwVaDQagwHDEXABtDgqaLXa6mq4oRQugP1FVGCz2Tod3IQEF8BWHBWUSqVcbu9Zj3ECtDiE4ECLo4KDgwNKWfEh7wu0OCrodDpU9/GBmA+0OCrQaDTYiuMEaHFU0Gq1sBXHCdDiEIIDLY4KDAaDw0Ex5AJiPtDiqFBdXV1VZdcZ5/ADtDiE4ECLowIMicAP0OKoAEMi8AO0OITgwJWGqMBisWBIBE6ArTgqqFQqmUyGtQoIgBaHEB9ocVSAIyr4AVocFeCICn6AFocQHGhxVIB5VPADtDgqwDwq+AFaHBXgSkP8AC2OCnClIX6AFocQHGhxVIAJ3/ADtDgqwIRv+AEObKECTPiGH2Arjgow4Rt+gK04KsDFtPgBtuKoABfT4gdocVSg0+lMJhNrFRBgv7sno0RsbKzJZDKZTEql0mQycblck8lkNBpPnz6NtTT7BfbFLUnLli2vXLlCIpGQtwqFwmg0tmvXDmtddg3sqFiShIQEZ2fn2kccHR1HjRqFnSIItLhFCQkJCQ4Ort338/f379y5M6ai7B1ocQszYcIEJycn5DWfzx89ejTWiuwdaHELExQUFBYWhrz28/ODTTjmQItbnoSEBCcnJz6fn5CQgLUWiB2PqOg0JnFhtVphsHjJTODbNnCAUqlsxg3PeWD5VeMMFsXFg05nwebJLOx0XPzifnFOpqKZD5PiYHtGoZBBwQtV8zbsz0a7Yq3FBrA/i5vAie1FXoGclhE8rKV8FAVZqgfXJENneFIdSFhrwTV2Z/GTO4r8QnnegUQIrCwv1Nw9Kx72nRfWQnCN7d2mP4b8LBWNSSWGvwEAQg+6iyfzRQYMEm0M+7K4pEhLoxPqkhkcivgNDC9qDEL93k2iUuh5LoTaDpPvTKtWGbFWgWvsy+IGPTDoCGUIg9GkrSbUFVkc+7I4xA6BFocQHGhxCMGBFocQHGhxCMGBFocQHGhxCMGBFocQHGhxCMGBFocQHGhxCMGBFscFBoPh0aNMrFUQE2hxXLBm3Y/rN67AWgUxgRa3DIVFBR8TP6XVaCwqB/Iv9huBbyZ/n005fvzgq9wcJpMV1a7DtKmzBQJHAIBOp9u955eLl/5Wq1WhoZHZ2c9Gj5oYO2goACAj897OXVtevsx2dHSKCG83ccJUZ2chAGBgbLeZMxakpl65fSeVzeYMHDAkYcxXAIBVqxOvXL0AAIjp0RYAkPzXKVfXZlhfN3GAFm+Cp08feXv79OrVTyqtOHosWalSrly+EQCwfcfPKSmHJ06YKhSKftm+QaOp7ttnEAAg/X7a/AXf9OrZb3DccIVcduTo/u9mT0765Q9kd6tVP/3f2IRJX3yRcPXqhb2/JQW0CmzfPnrUiPFl4tLi4sIF85cCAJycnM3QBTEXaPEm+O7bhTWZZqlU6h9/7tZoNFQq9dSpo/37xQ0fNhoAYDKZlq9Y/Ohx5ieRUZu3rBk4IP6b6XORj7Rt2z5h3NC79251jo4BAPTrGztyxDgAQAv/VqfPHE+7d6t9+2hPT28+X1AhlYSEhGN6rcQEWrwJdDrd0WPJFy6eEYtL6HSG0WisrJQ6ODhotVoPj7eh78gLhUJeUlKcl5dbWPjm1OljtQsRi0uRFwzG27z6FArFxUUkKS+z+gXZHdDijWEymRYumpmV/TRhzNdt2oRev345+cDvRpORzxdw2JxHjzI/HzoSAPDs2WMAgL9fS6lUAgBIGPN1l87da5fj5CR8t3AqhWowWj4XF6QO0OKN8eDB/fT7aYsWLuvZow8AoLAgHzlOoVC+/HLszl1bli1fJBSKTqQcGhL/pZdX8zdv8gAAGk21t7fP+9ZlbwltrAYcNGwMmbwSANCqZevab41GIwAgLnZYu7btpdKKqirFooXLpk2dBQDw9PR2dW3299kUtVqNfESv15uzASeDwayokCAlQywLtHhjtAkModFoO3dtuX3nxl/79+79LQkAkPsqBwDw4/KFPB6/X7+4iIh2JEAqLS0BAJBIpKn/myWRlE+dPvb4iUNHjyZPnTb2RMqhJisKC41UKOTrN6w4d+5U2t1bVrk4ewF2VBrDxUW0eNHyrdvWJS6ZG9QmdP26pD17tx89lhwd3S0yot3e35IuXT6HnEmhUObO/qF37/6do2NWLt+4Z+/2rdvWsdmc0JCI0NDIJivq1atfVvbT8xdO37p9PXbQ51HtOqB/cfaCfeU0vHa0nMGhBn4q+PiiDAYDhUJBXssV8vkLvqFSqZs27vr4kt+L10+rCrKq+o6FU0UNAlvxD2Td+uUvX2Z36NBFIHDMf/P61asX/fsPxloUpB6gxT+QqKiOYnHJkaN/6XQ6NzePMaO/QgYQIXgDWvwD6da1Z7euPbFWAWkaOKICITjQ4hCCAy0OITjQ4hCCAy0OITjQ4hCCAy0OITjQ4hCCAy0OITjQ4hCCY18T+Ew2mUIh1F81mUTi8O3rR3xfCPV7NwnX2aH0jRprFZZEnK/mCKDFG8O+LO4dwFYp9FirsCQKqa55IBtrFbjGvizO4pJDo/mX/yrGWohluHak1DeI5dTMAWshuMa+on4QXj9VXT9WFthe4OzGoDNt749cpzNKCjV5TxWt2/ICP+ViLQfv2GM3zqcNiy90f/BPZcHzqspybVVVFY/Hs2wVapXKYDRyOBzLFovgKKKx+JQ8xUVHdfNA0A2NKgiFyb6ZMmWK0Wi0bJk6nS4uLm7AgAFlZWWWLbkOs2bNQrV8YmB7t2lLkZ6eDgDYtm1bTcpCS3Hw4MGioqKioqJDh5pOL/ExrF27FgBw+fJlDczd3DB2avE//vjj9evXaJSs0+mOHDliMBhIJNL58+crKirQqKU2ERERMTExSqUS7YpsFDu1OJPJHDJkCBolJycnFxQUIK8LCgoOHDiARi21cXR0vHnzplQqLS0tRbsuW8TuLJ6UlAQAQMnfRqPx+PHjBsPbZJwmk+n8+fPl5eVo1FUHT09PEon01VdfWaEu28K+LL5s2bL+/fujV/6BAwdqmnCEwsLCgwcPoldjbUQi0ZQpU06dOlVVVWWdGm0CexkXz8/P9/b2lsvlFh8frE18fHxeXh6JREIScJJIJBKJ5O7unpKSgl6l75KTk5Oamjp27FhrVopb7MLiz58/T0pK2rBhg9Vq3Ldvn0QimTlzptVqrMPmzZv79u3bokULrATgB7voqNy/f9+a/gYA0Gg0Go1mzRrrMH36dKFQWFhYiNLAkQ1BcIsfPnwYADBixAgr16vRaMxJK44qAoHAzc1t9uzZjx8/xlYJthDZ4lu3bnVwwGaJkslksviM0gdAJpMPHz5c5wnY3iCyxQMCAmJjYzGp2sHBgc3GyxrXPn36AABWrVqFtRBsIKbFkZntnj0xS6upUqkw76jUYfjw4d9++y3WKjCAgBb/+eef4+LisFYBsH3cfBdfX981a9YAAO7evYu1FqtCQIt//vnnmA+WVVVVYfUY0AhUKhUZXzpx4gTWWqwHoSy+ePFisVjs7u6OtRCgVCrx0xevw6RJk/DwKGw1iGPx5cuXz5w5UyQSYS0E4NziAIBBgwYBALZs2YK1EGtAHIsvWrRIKKxnk2JMMBgMfD4faxVNEBsb++WXX2KtAnWIYPHFixe/fPkSaxX/4c2bN46OjliraAIvL6/t27cDAORyOdZaUMTmLb5r166EhAR/f3+shfyHiooKZ2dnrFU0DXKr2bdvX05ODtZa0MLmLT5x4sSWLVtiraIunp6eTk5OWKswl6lTp27duhVrFWhhwxZPTk4+duwY1irqoaSkpLi4mEy2pe8WWab25MkTrIVYHlv6GWqTlpZGJpMHD8bjbq5FRUV4GLj8ANLT0+/du4e1CgtjqxaPiooaNmwY1irqp7Cw0MPDA2sVH8KYMWOIZ/EGUwXhNjjKZDKdO3cOWVr0AZ/lclFPHyWXy1u3bo12LSgxefLk8vJyMplstQUIbDYb1amoBi2uUqnQq/VjUCgUnTp1+jB5JBLJCha/e/cuSuHP1kEoFGZkZLi4uFjH5SjlDKvB9joqXC6XQqFgraIxsrKyAgICsFbxUXh7e5PJZGIEPdqSxQ0GQ3V1NdYqmqCyslKn0+FkHcHHQKVSkSBrW8eWLF5ZWUmn07FW0QQ5OTndu3fHWoVloFAoEonE1ttyzCz+/PnzOpn41q9fP2PGjIbON5lMzs7O+F8id/v2bRsdMawXZ2dnS8V2FBUV9evX7+rVqxYpzXywsfiFCxe+++67Or0OFovFZDLrPd9kMtnKTfPevXtt27bFWoUlodFoNt2QY5NfXKvVvntw8uTJDZ0vk8nQfu62CNXV1Tk5OcHBwVgLsTAajUav19vET/Au72Hx6urq5OTkf/75RyKRiESiHj16DBs2jEKhVFRU7Ny58969ewaDoU2bNhMmTPD19QUALF261NPTk0KhnD17Vq/Xt2vXburUqWw2+8KFC8iKCGQl57fffturV6+xY8eKxeI2bdogYZeff/751KlTb926lZaWxmaze/fuPWbMGABARkbGokWL1q9fXzPwPHjw4EGDBo0bNw6ZOd+5c2dGRgadTvf39x8zZkyrVq1Q++rqITMzs2/fvtas0TowGAytVnvy5Mnjx49LJBJXV9du3brFx8fT6fSXL1/Onj17yZIle/bsyc3NFYlE48ePb9++PfLBysrKHTt23L59m06nh4aGYiLe3I6KwWBITEw8evRop06dZs6cGR0dXVBQQKFQqqurFyxYkJmZOX78+GnTpkkkkoULF9ZMGx09erS0tDQxMXHSpEmpqanJyckAgLZt28bHxwMAEhMT16xZg9zWv/nmmzqrBdevX+/n57d69eru3bsnJyenpaU1rrCiomL27NkKhWLSpEnjxo3T6/Vz5861cqKcCxcuBAUFWbNGq3Ho0KG9e/d26dJlxowZ0dHRhw8f3rx5M/JfGo1m5cqVcXFxq1atEolEq1evlslkyL160aJFt2/fHjx48Lhx40pKSjBRbm4rnpqa+vDhwxkzZnz22We1j1+5cuXNmzcrVqwIDw8HAAQFBY0fPz4lJQXJzuPh4TFnzhwSiRQQEHDjxo309PQJEyY4Ojq6ubkhSSBq4gYiIyOPHj1au3feu3fv4cOH6/V6Nze3c+fO3b9/PyoqqhGF+/fvFwgEK1asQCIUu3fvPnHixHPnzk2aNOmDvpkP4dKlSxgmeUMPiURy4MCBOXPmhISECAQC5DF0y5YtNd/t5MmTu3btCgAYO3bsN9988/jx406dOp06dSo3N3f58uUREREAgMDAQGv+FjWYa/H09HQ6nf5u2oaHDx+y2WzE3wAAV1dXLy+v7Oxs5C2dTq8ZA3F1dX327Jn5yhgMBgBArVYzmUxnZ2eJRNL4+ffu3SsrK6s9rajT6crKysyv8SNJS0sLDAy0wuyp9cnIyNDr9UgfEgF5AK35UZAfC8mOW3P85s2bPj4+iL+RIUgstJttcalU6uTk9K5KlUpVJ4KLy+XWuzUClUqtSbxtPiwWi0KhmPNZqVQaFRWFdMprsGYA5aVLl3r06GG16qwJ8oMmJibWCR10c3PLy8urfQTJO4AMf5WVleEhVMVci3M4HKlU+u5xZ2fn58+f1z4ilUpdXFzMKdOcoag6f1SNjItzOBy5XO7l5WVO1Whw/vx5oubiqbk1eXl5GY1GtVptTtvB5/MrKyvRV9cE5j5uhoWFVVdX1x631+v1SAdLoVDUuDw3N7eoqKjJRy7kvtbkPjhyuRyppQakI1hzf6yoqKg5ITw8/OnTpy9evKg5Wa223l7gp0+f7ty5c839mmCEhYWRSCQkSzqydsUc7/r7+7948QLzjIrmtuIxMTEnT55cv359dna2n5/f69evMzIyNm/eHBMTc/DgwZUrV3755ZckEik5OZnP5ze5E0ObNm0oFEpSUlKvXr20Wm2/fv3qPU2r1dbJeO/p6SkSiZKTkwUCgVqt/u2332qmhEaOHHn37t3FixcPHjxYIBCkp6cbDIYffvjBzAv8SA4fPkzUJhwA4O7uPmjQoBMnTiQmJnbo0EEqlZ48eXLJkiWNp2QaNmzY5cuX586dGxcX5+TkZP15TQRzW3E6nb5y5coePXpcuXJl27Zt6enp0dHRer2eSqUuW7asZcuWO3fuTEpK8vT0XL16dZPB525ubtOnTy8oKEhKSrp27VpDp72bNIJKpS5cuJBKpS5evHj37t0jRoyoWfDp5ua2du3awMDAgwcP7tixQyaTxcTEmHl1H0l2dnZ1dTVW477W4euvv544cWJeXt7WrVvPnj3bsWPHJuOv3dzcli5dKhQK//zzz/379yOzJdanwV0ixGKx1cWgDolEMvM54b345Zdfmjdv3tC9yBaRSCRNPt+Xl5dbJHEN2qsy8bvSUKVS4TYsozZSqfTIkSNE8reZcLlcm9jSFps1KuZgNBpxmPnyXX799dcJEyZgrQID8L+wGQG/FreJRT9qtfr48eOpqalYC8EGrVZLoVBwHoSF346KTbB///6pU6dirQIzSCQSbsPYa8BvK65QKOh0Ot4S0ddGpVLt2bPn+vXrWAvBDAcHBxaLhZONjRoCvxY3c/oTQzZt2tRImJKdgP/npQYHDetMK1ofjUaDrE6xbLGWKrCoqGjSpEknT560SGl4w2g0mhlmVVJScuzYsSlTpnxwXRb/ietgF7sno8H8+fN79uyJ4Y5Z+KFPnz779u1DY8LBIuD3cfPq1atWm35/Xx4+fCiTyaC/EZAk5bgFv31xHx8f3G5uvXLlyiVLlmCtAi/4+PhgLaEx8NuK+/j4/P7771irqIcjR46EhIRYOSoUz8hksmnTpmGtokHw24ojAaM4nFZYuXKlvW1d2Th8Pj8/Px+3+Xhx/bi5Y8cOk8mEScBfQ/z888++vr7IjmeQGmQyGZ1Ox+dyefx2VAAAnTp1ysjIwFrFvzx69CgjIwP6+114PB5uJ+lw3Yrjjfj4+A0bNjRv3hxrIbjj6tWrJ0+eXLduHdZC6gHXrTgAoLS0FA/xfwCAPXv2dO/eHfq7Xnx8fOrEKeMHG7A4HgLGSktLjx49iudxA2zx8fHZtm0b1irqB+8WDw0NDQgIsGY6lHpZsGDBsmXLsNWAc3CbUh3vFkemyrGdHE5OTg4MDAwLC8NQA/6ZN29eVlYW1irqwQYsbjAYsAreRhJP3rp1a86cOVgJsBV0Oh1WWQsbB9dTPwgUCiU1NbWysjIuLs76tc+aNWv69OnWr9fmWLZsGdprBj8MG2jFAQDTp0/HZFph3759ISEhNRkbIY2Aw3loBDgu3iCFhYUbNmyonasS0gg7d+40GAyN7IOAFbbRigMAHj9+PG/ePABAr169IiMjrTDL8O23337MSn97g0aj1bv5B+bgsfNUL8HBwVeuXImMjCSTySQSCe0MB1u2bOnbty8eEqvaCnFxcZba+Mqy2IDFY2NjKyoqlEolmUwmk9/edupkfLYsT548yc3Nxed0NG5B9Rf5GGygoxIREUGn02vMjYBqpvoZM2YsXrwYvfIJyaVLl3799VesVdSDDVg8MTFx5MiRyN4pCBQKpY7jLcjSpUunT5/eZOZRSB1UKlWT2bQxwQY6KgCAcePG+fv7b926NScnh0QiUalUlHIbXL16VS6Xx8bGolE4sRk4cODAgQOxVlEPtmFxAECXLl38/PwWLlz47NkzCoWCxuOmwWCYO3duk1vDQepFq9Xq9XoWi4W1kLrgwuJVUoPB0HTWDi7DdfP6X9esWfP48WOTlikrt/Dz+9KlS1f9uKmhYikUMscRp7MbeODMmTOPHj36/vvvsRZSF4ynfv45Up6dLhd5MyvF7zGkatDrKZaeKzaZTACYSKQGu/gCEU2cr271Ca/rEAsk1SYM/fv3Ly0tRbILkUgkJPObk5PT+fPnsZb2FsxacYPO9MfKvLa9XYI6OdGZNvDUCwDQqI2leerflr4etaA5xQG/WfysyfDhw7du3WowGGoPAHzyySeYivoPmHnrz5/yY75w9w5k24q/AQB0Jtm7NTvmS/c/V+VjrQUvDBkypM4uec2aNRs5ciR2iuqCjb0yr1a2jhI4uuI0oLVxHEW0wPaCjCu4CLfDHDab3b9//9prsMLCwoKDgzEV9R+wsXjBCzW5pyUSAAAK/ElEQVRHgIsn3Q+DzacW5Fhvx0OcM3To0JqQVldXV1w14Rh2VEiOItvYRqNeHEV0EoB98bdwudy+ffsi+3GGh4e3adMGa0X/ARuLS8Uaoy0v4jWaTNJSG9jJyWoMGzbM39/f3d19+PDhWGupiw33FiAfjKxcV5qnqZTolDKDCQC1wgK55GOC5ijk8jf3RG/ufWx4G51JIZEAm0/lCCgiT7rI66Nu+NDidoRcon98S/Yis0qnMTEFTAqVTKVRaEwHI7DAagiRu7/IHVhkNs5QTdJpDeVig06jNerkWpXOL4QT2I7r5vchkV/Q4naBttr4zxFJ0atqtpDVLMCVzsH77iW10WsMlWJl6qlKCsXYNV7o7PZ+A3HQ4sTn0Q1F6gmxWyvn5m2dsNbyIVDpFCcvHgBAUa5K2VHcMoIbPeg9LsRmpl0gH8bF5LJn99SBMT4CDxRX2FsHrpDlG+UpKScf2lho/qegxYnMxf3lcjlV1MoZayGWhO/GZboIfl+ebzJrvy1oceKSsqNYJn97iycYHCeGqIXLniVm7ZMDLU5Mbp6u0GgdnL0J6G8EBo8m9HM+tq24yTOhxQnImyx1cZ7exZ/gsXk8EYtEYzS5WAhanIBcPVLGdiFs+10bgQfv5qlyg76xmXJocaKRdU9BZdAYNjXy/TG4BTinnihv5ASbsbjBYHj0KPMjC/l500/xQ3tbSBFOeXKnSuiLx/Hvcsmb2d9/mvHQwtFATl68kjxdtbLB4RWbsfiadT+u37gCaxV4R1KslZXrHBj2FWNqIlNePa5q6H9txuJaDVzZ1zSvHlWxnXAXA482LEdWTqayof+1jQn8VasTr1y9AACI6dEWAPDXnyluzdz1ev2evdvPnT8lk1U2b+47NmFSdKduyPlPnz3enrQxK+spg8Hs2KHLlCnf8rj1PH79tX/v8RMHFQp5ixYBYxMmfRIZZfUrszCl+TquCK0HzZtpR/658ZdMLnZydI8I7d2t0ygHB3phUdaWXV9NGL3hzPltRSXZjgK3/r2nBQd2QT5SpZSeOLPhyfNrDlS6vy9aAZ08EavosRyYQL1r+G3D4qNGjC8TlxYXFy6YvxQA4OwkBACsXbfs4qW/R40c7+Pjf/HS39//MPvnDTtDQyNev341a/ZkHx//uXP+T1Yp3bN3u1hcsm7tL3XKTL+ftnPXlh49+nzarmPa3ZtqlQqji7Mkxbkq33aojBWev7zznxt/RXcY7uriKy7Pu3r9j/LyN18OTQQA6HSaPw4sius/y1Hgdu7yjr8Ofb9o1gk2W6DTa5P2TpdI3nTpNNLJ0e3mnSNoCENQVmqrZPp6Q8lsw+Kent58vqBCKgkJeZvNPj//9bnzp8aMnjg2YRIAoGuXHqPGDN77W9L6ddv/+PNXMpm8+qctXA4XAMDl8las+uHBg/thYZG1yywpKQIADI4dFhQU2qtXP4yuzMJoVAYq3fIdcZm87NK1vSOH/hga3B05wucKj5z8Kbbfd8jbuP6zwkN6AQD69frfxl8SXr7OCA2KuXH7UHHJi68TNrdqEQUA8PEKWb0JrYAJBwZVJTfYsMXf5cHD+wCA6OgY5C2JRGrXtv2Fi2cAAJkP0iMi2iH+BgC0a9cBAJCV/bSOxdt/Gs3l8las/H76tDnt20djcREWRqUwMDmo/KAvXqYZDPo/D//w5+Ef/v8xEwBAphAjb2gOTOSFo8ANACBXlAEAHj/7x821BeJvAACZjOJDMI1BVSkM9f6XrVpcqawCADgK/h0d4/H4KpVKqVQqlVUC/r83ay6XBwAoL6+7raGzs3DLpt1bf1m/YNHM4OCwHxavdHHB6bZ6ZkIiAb3evKVJ74lcUQ4AmDBqvYD/n6/I2cmzpPRl7SNUigMAwGg0AAAqZSUebgFo6HkXg9FIaiCY1mZGVP5/wqq3CIUiAIBcLqs5UlEhoVKpDAZDKBTVPi6VVgAAOJx6lpJ6e/v8tHLTurW/5Obm/LQ6Ef0rQBcmh6LTGAEKMbFM5ttHWJGLT+1/FEpjTSSH7VillFpeTX0YNAYWr/67hM1YnMFgVlRIjMa3rVRgYDCJRLp9JxV5q9Vqb99JDQoKpVAoQUGhmQ/Sq6urkf+6du0SAADpxDs40NRqlV6vr/kUACAyol379p2zXzzH6MosCYNN0Wnqv19/DC392pJIpNQ7B2uOaLRN59jwcAt4U/hUXGaNjcO11QYWt/6/N5vpqISFRv59NmX9hhUhweFcLq9jxy6f9R6w97ckg8Hg7u55+vSxigrJwgU/IsMvly+fm7dg+sABQ8Tikt9+3xER3jY87BMAQMsWAdXV1YlL502Z/K1cLluydF5c7DAmk5WWdrN1AL5SI3wY7n4sXbXe4lM/Qmev6PbDr99K3v3HrKDArgpF+Y07hyeMXu/p3rqRT8V0HnMv88y23ZO7dPiCxxXef3jOsqpqMBkBR0Bl8+u/apuxeK9e/bKyn56/cPrW7et9PhvYsWOXmTPms9mcY8cPKBRyXx//Fcs2REa0Q4ZfVq/asmPX5tVrljCZrF49+02eNBNJJ9mjR5+cl9mXLp99nfuyWTP35t6+f/21x2QyhYV/8s20uVhfogVw86W9eFTFElg+R82gvjMFfFHq7UNZObd5XGFwm258XhOPLkJnz6/G/Hzq3KZzl3cK+K4hgd2yc+5YXBgAQC5WCoQNrsnBJjPtvuV53Ue485xsdamQvEJ3+c+i0YubYy2kLrJy3eFNhf4dvMw4lzgUPRW3jeG0iqw/cs9mWnGIOfCFDi4eDI1ST2c3+Mv+nrww+2U9ramA51opL333OJvJX/DdUQuK3LprUnFpzrvHPd1aFxTX/0S0ZP65Rh5tycDoF9pgZCq0ONEIiebeOC3xDHFt6IS4/rP0+noW/Oj1Oiq1nvtqIznXP4xRw5YZDPUkXCGRGuxTNDKmXp5b2TyA0Ui6eWhxouEbxL5zVqqq1DTUI+dxMY5W5vNcLFWUyWgSv6wc+r/Gtke1mUFDiPl0H+ZSXaHAWoU1kBXJun7exB8MtDgBEXnRW0UwxC8ai4UhALJiBYdtCGrfxMpKaHFiEtyRL3Qli3PwuBGmRZCVKNUVip4jml5zAS1OWLoPd/ELpJW/IqDLZSVVRnXVF7PNGhuFFicybXvym7ekFj8pNTYao25bVORX0snVg//nbub50OIEp11vx04DHLOv55W9stKKKPSQFiieX3nt15raJ6HBIdF3gYOGxMe7NXPyav+0c9IH1/K5QjbHhcVxZmIt6j1QyzSKMpVRq3XxcOi/1JfGeL92GVrcXoj6zLFtT8cnt2TZGfL8zFJBMxYgAQqN6sBwMKKzyvyDIVNIeq3BoNXrtQa9xkBnklqGcwI+EfGcP8Su0OJ2BJkCQqL5IdF8g85UkletlOmVcoPBYNKo8NVTd6ABCpXC4tHYfKrQjd7QQnAzgRa3RygOJI8WttRX+Riwedx0akYjNxSHZAuQSSSnZja8qaJdgY3FyWSSpLgak6otQkWJBpDwdXOHNAQ2FvcOYCqkFtncCxsUUp1XK7tLOmWjYGPxoI78N1nKvKcN5qHDM/nPlPnPFKGd+VgLgZgFNlE/AACTCRzZVOgbwhV5MQSi99tmDisqxVpxfnXuY/mQbzxt+VHCvsDM4gj3Lkiz7yvoTHJ5Id6zcgo96Bq1sVUEt21vgu++QDAwtjiCwQDwv4iCTCVR7CunMUHAhcUhEPSAy7AgBAdaHEJwoMUhBAdaHEJwoMUhBAdaHEJw/h8Sp1Bl6ut1hwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with suppress(Exception):\n",
    "    display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Agent\n",
    "\n",
    "We generate a unique user name and thread id (session id) and add these to Zep, associating the Session with the new User."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_name = \"Daniel\"\n",
    "last_name = \"Chalef\"\n",
    "user_name = first_name + uuid.uuid4().hex[:4]\n",
    "thread_id = uuid.uuid4().hex\n",
    "\n",
    "await zep.user.add(user_id=user_name, first_name=first_name, last_name=last_name)\n",
    "await zep.memory.add_session(session_id=thread_id, user_id=user_name)\n",
    "\n",
    "\n",
    "def extract_messages(result):\n",
    "    output = \"\"\n",
    "    for message in result[\"messages\"]:\n",
    "        if isinstance(message, AIMessage):\n",
    "            role = \"assistant\"\n",
    "        else:\n",
    "            role = result[\"user_name\"]\n",
    "        output += f\"{role}: {message.content}\\n\"\n",
    "    return output.strip()\n",
    "\n",
    "\n",
    "async def graph_invoke(\n",
    "    message: str,\n",
    "    first_name: str,\n",
    "    last_name: str,\n",
    "    thread_id: str,\n",
    "    ai_response_only: bool = True,\n",
    "):\n",
    "    r = await graph.ainvoke(\n",
    "        {\n",
    "            \"messages\": [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": message,\n",
    "                }\n",
    "            ],\n",
    "            \"first_name\": first_name,\n",
    "            \"last_name\": last_name,\n",
    "            \"session_id\": thread_id,\n",
    "        },\n",
    "        config={\"configurable\": {\"thread_id\": thread_id}},\n",
    "    )\n",
    "\n",
    "    if ai_response_only:\n",
    "        return r[\"messages\"][-1].content\n",
    "    else:\n",
    "        return extract_messages(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello! How are you feeling today? I'm here to listen and support you.\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"Hi there?\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm sorry to hear that you've been feeling stressed, especially with work and your dog. It can be tough when multiple things weigh on us. Would you like to share more about what's been causing the stress at work or any specific concerns you have about your dog? I'm here to help.\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"I'm fine. But have been a bit stressful lately. Mostly work related. But also my dog.\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It's completely understandable to feel worried when a beloved pet is sick. Our pets are like family, and their health can deeply affect us. What symptoms is she showing, and have you been able to take her to the vet? I'm here to support you through this.\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"She's sick. I'm worried about her.\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wait a few seconds to allow Zep to ingest the new messages\n",
    "\n",
    "Zep takes a few seconds to ingest messages and build the user knowledge graph. The last few messages in the chat thread being passed to the agent should ensure the agent remains grounded in the recent conversation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Viewing The Context Value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "FACTS and ENTITIES represent relevant context to the current conversation.\n",
      "\n",
      "# These are the most relevant facts and their valid date ranges\n",
      "# format: FACT (Date range: from - to)\n",
      "<FACTS>\n",
      "  - Daniel99db is worried about his sick dog. (2025-01-24 02:11:54 - present)\n",
      "  - Daniel Chalef is worried about his sick dog. (2025-01-24 02:11:54 - present)\n",
      "  - The assistant asks how the user is feeling. (2025-01-24 02:11:51 - present)\n",
      "  - Daniel99db has been a bit stressful lately due to his dog. (2025-01-24 02:11:53 - present)\n",
      "  - Daniel99db has been a bit stressful lately due to work. (2025-01-24 02:11:53 - present)\n",
      "  - Daniel99db is a user. (2025-01-24 02:11:51 - present)\n",
      "  - user has the id of Daniel99db (2025-01-24 02:11:50 - present)\n",
      "  - user has the name of Daniel Chalef (2025-01-24 02:11:50 - present)\n",
      "</FACTS>\n",
      "\n",
      "# These are the most relevant entities\n",
      "# ENTITY_NAME: entity summary\n",
      "<ENTITIES>\n",
      "  - worried: Daniel Chalef (Daniel99db) is feeling stressed lately, primarily due to work-related issues and concerns about his sick dog, which has made him worried.\n",
      "  - Daniel99db: Daniel99db, or Daniel Chalef, is currently experiencing stress primarily due to work-related issues and concerns about his sick dog. Despite these challenges, he has shown a desire for interaction by initiating conversations, indicating his openness to communication.\n",
      "  - sick: Daniel Chalef, also known as Daniel99db, is feeling stressed lately, primarily due to work-related issues and concerns about his sick dog. He expresses worry about his dog's health.\n",
      "  - Daniel Chalef: Daniel Chalef, also known as Daniel99db, has been experiencing stress recently, primarily related to work issues and concerns about his sick dog. Despite this stress, he has been feeling generally well and has expressed a desire to connect with others, as indicated by his friendly greeting, \"Hi there?\".\n",
      "  - dog: Daniel99db, also known as Daniel Chalef, mentioned that he has been feeling a bit stressed lately, which is related to both work and his dog.\n",
      "  - work: Daniel Chalef, also known as Daniel99db, has been experiencing stress lately, primarily related to work.\n",
      "  - feeling: The assistant initiates a conversation by asking how the user is feeling today, indicating a willingness to listen and provide support.\n",
      "</ENTITIES>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "memory = await zep.memory.get(session_id=thread_id)\n",
    "\n",
    "print(memory.context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That sounds really frustrating, especially when you care so much about your belongings and your dog’s health. It’s tough when pets get into things they shouldn’t, and it can add to your stress. How are you feeling about that situation? Are you able to focus on her health despite the shoe incident?\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"She ate my shoes which were expensive.\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now test whether the Agent is correctly grounded with facts from the prior conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We were discussing your concerns about your dog being sick and the situation with her eating your expensive shoes. It sounds like you’re dealing with a lot right now, and I want to make sure we’re addressing what’s on your mind. If there’s something else you’d like to talk about or if you want to share more about your dog, I’m here to listen.\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"What are we talking about?\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's go even further back to determine whether context is kept by referencing a user message that is not currently in the Agent State. Zep will retrieve Facts related to the user's job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You've mentioned that you've been feeling a bit stressed lately, primarily due to work-related issues. If you'd like to share more about what's been going on at work or how it's affecting you, I'm here to listen and support you.\n"
     ]
    }
   ],
   "source": [
    "r = await graph_invoke(\n",
    "    \"What have I said about my job?\",\n",
    "    first_name,\n",
    "    last_name,\n",
    "    thread_id,\n",
    ")\n",
    "\n",
    "print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
